import cv2
import numpy as np


class yaw(object):
    def __init__(self):
        self.medianBlur = None
        self.Cannymin = None
        self.Cannymax = None
        self.Threshold = None

    def image_init(self,medianBlur,Cannymin,Cannymax):
        self.medianBlur = medianBlur
        self.Cannymin = Cannymin
        self.Cannymax = Cannymax
    def parameter_init(self,Threshold):
        self.Threshold = Threshold
    def get(self,frame):
        global angle
        frame1 = frame.copy()

        GrayImage = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        GrayImage = cv2.medianBlur(GrayImage, self.medianBlur)
        edges = cv2.Canny(GrayImage, self.Cannymin, self.Cannymax, apertureSize=3)
        
        try:
            lines = cv2.HoughLines(edges, 1, np.pi / 1800, self.Threshold, min_theta=1.22,max_theta=1.91)
            lines = np.reshape(lines, (-1, 2))
            lines = lines.tolist()
            lines = sorted(lines, key=(lambda x: x[0]), reverse=True)
            current_rho = lines[0][0]  # 当前距离摄像头最近的识别到的线
            current_theta = lines[0][1]  # 当前距离摄像头最近的识别到的线
            for i in lines:
                angleSign = False
                rho = i[0]  # 第一个元素是距离rho
                theta = i[1]  # 第二个元素是角度theta
                if abs(current_rho - rho) < 10:
                    current_rho = rho
                    current_theta = theta
                else:
                    pt1 = (0, int(current_rho / np.sin(current_theta)))  # 该直线与第一列的交点
                    pt2 = (frame.shape[1], int((current_rho - frame.shape[1] * np.cos(current_theta)) / np.sin(current_theta)))  # 该直线与最后一列的交点
                    cv2.line(frame1, pt1, pt2, (0, 0, 255), 1)  # 绘制一条直线
                    angle = (current_theta / 3.14) * 180 - 90
                    angleSign = True
                pt1 = (0, int(rho / np.sin(theta)))  # 该直线与第一列的交点
                pt2 = (frame.shape[1], int((rho - frame.shape[1] * np.cos(theta)) / np.sin(theta)))  # 该直线与最后一列的交点
                cv2.line(frame1, pt1, pt2, (255, 255, 0), 1)  # 绘制一条直线
            if angleSign == False:
                angle = (current_theta / 3.14) * 180 - 90
            if len(lines) == 1:
                angle = (current_theta / 3.14) * 180 - 90
            return angle,frame1
        except Exception as e:
            return None, frame1